import numpy as np
from multiprocessing import Pool
import DRW_library as dl
import time
import celerite
from celerite import terms
from scipy.optimize import minimize
start=time.perf_counter()

def obs(period):
    if period==1:
        t=np.arange(1,181)+np.random.uniform(-0.13,0.13,180)
        return t
    else:
        t=np.arange(1,181)+np.random.uniform(-0.13,0.13,180)
        for i in range(5):
            t1=np.arange(1,181)+np.random.uniform(-0.13,0.13,180)
            t0=(i+1)*365.25+t1
            t=np.append(t,t0)
        return t

def mainprocessing(m):
    tau=0.1*1.2589**m
    at=np.array([])
    for k in range(N):
        np.random.seed(seed[round(N*m+k)])
        #模拟在某天的不同时刻观测
        t=obs(period)
        y=dl.DRW_process(t,tau,0.2,18)
        e=np.full_like(y,sigma)
        s=np.random.normal(y,e)
        try:
            re=dl.DRW_fit(t,s,e)
            at=np.append(at,re[0]/tau)
        except:
            print(tau,'Err')
            at=np.append(at,-1)
            continue
    print('[',m,']',tau)
    return np.concatenate(([tau],at))

#BEGIN
N=10000
period=6
sigma=0.014
seed=np.random.randint(0,2**31-1,N*100)
print('N=',N,'sigma=',sigma,'period=',period)
if __name__ == '__main__':
    pool=Pool(21)
    l=np.array(pool.map(mainprocessing, range(21)))
    pool.close()
    pool.join()    
np.savetxt('1t6e1.txt',l,fmt='%.4f')
print('time=',time.perf_counter()-start)
#END
